Controlled Sparsity via Constrained Optimization or: How I Learned to Stop Tuning Penalties and Love Constraints. (arXiv:2208.04425v2 [cs.LG] UPDATED)
The performance of trained neural networks is robust to harsh levels of
pruning. Coupled with the ever-growing size of deep learning models, this
observation has motivated extensive research on learning sparse models. In this
work, we focus on the task of controlling the level of sparsity when performing
sparse learning. Existing methods based on sparsity-inducing penalties involve
expensive trial-and-error tuning of the penalty factor, thus lacking direct
control of the resulting model sparsity. In response, we adopt a constrained
formulation: using the gate mechanism proposed by Louizos et al. (2018), we
formulate a constrained optimization problem where sparsification is guided by
the training objective and the desired sparsity target in an end-to-end
fashion. Experiments on CIFAR-{10, 100}, TinyImageNet, and ImageNet using
WideResNet and ResNet{18, 50} models validate the effectiveness of our proposal
and demonstrate that we can reliably achieve pre-determined sparsity targets
without compromising on predictive performance.